import os
import json
import argparse
import time
from typing import Optional, Dict

from loguru import logger
from elasticsearch import Elasticsearch
from elasticsearch_dsl import Index


def load_mappings(mappings_path: Optional[str]) -> Dict:
    if mappings_path is None:
        repo_root = os.path.abspath(os.path.join(os.path.dirname(__file__), os.pardir))
        mappings_path = os.path.join(repo_root, "configs", "json_files", "mappings.json")
    with open(mappings_path, "r", encoding="utf-8") as f:
        return json.load(f)


def get_es_client() -> Elasticsearch:
    es_url = os.getenv("ES_URL", "http://localhost:9200")
    es_user = os.getenv("ES_USER", "elastic")
    es_password = os.getenv("ES_PASSWORD", "8899aabbccdd")
    return Elasticsearch(es_url, basic_auth=(es_user, es_password))


def index_exists(es: Elasticsearch, index_name: str) -> bool:
    try:
        return es.indices.exists(index=index_name)
    except Exception as e:
        logger.error(f"检查索引是否存在失败: {e}")
        return False


def create_index(es: Elasticsearch, index_name: str, mapping: Dict) -> bool:
    try:
        es.indices.create(index=index_name, settings=mapping["settings"], mappings=mapping["mappings"])  # noqa
        logger.info(f"索引创建成功: {index_name}")
        return True
    except Exception as e:
        logger.error(f"创建索引失败 {index_name}: {e}")
        return False


def delete_index(es: Elasticsearch, index_name: str) -> bool:
    try:
        es.indices.delete(index=index_name, ignore_unavailable=True)
        logger.info(f"索引删除成功: {index_name}")
        return True
    except Exception as e:
        logger.error(f"删除索引失败 {index_name}: {e}")
        return False


def ensure_index(es: Elasticsearch, index_name: str, mapping_key: Optional[str], mappings_obj: Dict, recreate: bool) -> None:
    # 推导 mapping key：传入优先；否则用 index_name 去掉最后一个 '_' 段
    if not mapping_key:
        mapping_key = index_name.rsplit('_', 1)[0]

    mapping = mappings_obj.get(mapping_key)
    if not mapping:
        raise ValueError(f"未在 mappings.json 中找到 key: {mapping_key}")

    if recreate and index_exists(es, index_name):
        delete_index(es, index_name)
        # 等待ES回收
        time.sleep(0.5)

    if not index_exists(es, index_name):
        ok = create_index(es, index_name, mapping)
        if not ok:
            raise RuntimeError("索引创建失败")
    else:
        logger.info(f"索引已存在，跳过创建: {index_name}")


def parse_args():
    parser = argparse.ArgumentParser(description="初始化或重建 ES 索引")
    # 支持无参数运行：从环境变量或默认值获取
    default_index = os.getenv("ES_INDEX", "yoka_kb_index_tang")
    default_mapping_key = os.getenv("ES_MAPPING_KEY", None)
    default_mappings_path = os.getenv("ES_MAPPINGS_PATH", None)
    default_recreate_env = os.getenv("ES_RECREATE", "false").lower()
    default_recreate = default_recreate_env in {"1", "true", "yes"}

    parser.add_argument("--index", default=default_index, help="索引名，如 yoka_kb_index_tang")
    parser.add_argument("--mapping-key", default=default_mapping_key, help="mappings.json 中的 key，不传则用 index 前缀")
    parser.add_argument("--mappings-path", default=default_mappings_path, help="mappings.json 路径，默认读取项目内文件")
    parser.add_argument("--recreate", action="store_true", default=default_recreate, help="如存在则删除并重建")
    return parser.parse_args()


def main():
    args = parse_args()
    mappings_obj = load_mappings(args.mappings_path)
    es = get_es_client()
    ensure_index(es, args.index, args.mapping_key, mappings_obj, recreate=args.recreate)


if __name__ == "__main__":
    main()

# python init-vb.py --index yoka_kb_index_tang --mapping-key yoka_kb_index_tang --mappings-path ./mappings.json
